import asyncio
import json
import os
from typing import AsyncGenerator
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.responses import HTMLResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from openai import AsyncOpenAI
from dotenv import load_dotenv
import uvicorn
from pydantic import BaseModel
import sys
import os
# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
from src.chatbot.converters.document_converter import convert_to_docx, convert_to_pdf, get_converter_status, is_pdf_available

# 加载环境变量
load_dotenv()

app = FastAPI(title="Streaming Chatbot")

# 初始化OpenAI客户端
client = AsyncOpenAI(
    api_key=os.getenv("OPENAI_API_KEY"),
    base_url=os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
)

class ChatBot:
    def __init__(self):
        self.model = os.getenv("MODEL_NAME", "gpt-3.5-turbo")
        self.max_tokens = int(os.getenv("MAX_TOKENS", "32768"))
        
    async def stream_response(self, message: str, history: list = None) -> AsyncGenerator[str, None]:
        """生成流式响应"""
        if history is None:
            history = []
            
        # 构建消息历史
        messages = [
            {"role": "system", "content": "你是一个友好的AI助手，请用中文回答问题。"}
        ]
        
        # 添加历史对话
        for msg in history[-10:]:  # 只保留最近10轮对话
            messages.append(msg)
            
        # 添加当前用户消息
        messages.append({"role": "user", "content": message})
        
        try:
            # 创建流式响应
            stream = await client.chat.completions.create(
                model=self.model,
                messages=messages,
                stream=True,
                temperature=0.7,
                max_tokens=self.max_tokens
            )
            
            async for chunk in stream:
                if chunk.choices[0].delta.content is not None:
                    yield chunk.choices[0].delta.content
                    
        except Exception as e:
            yield f"抱歉，发生了错误: {str(e)}"

# 创建聊天机器人实例
chatbot = ChatBot()

# 存储WebSocket连接
class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []
        self.chat_histories: dict[WebSocket, list] = {}

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)
        self.chat_histories[websocket] = []

    def disconnect(self, websocket: WebSocket):
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)
        if websocket in self.chat_histories:
            del self.chat_histories[websocket]

manager = ConnectionManager()

# 请求模型
class ConvertRequest(BaseModel):
    markdown_text: str
    title: str = None
    user_question: str = None

@app.post("/convert-to-docx")
async def convert_docx(request: ConvertRequest):
    """将Markdown文本转换为DOCX文档"""
    try:
        file_path = convert_to_docx(request.markdown_text, request.title, request.user_question)
        
        # 返回文件下载链接
        filename = os.path.basename(file_path)
        status = get_converter_status()
        return {
            "success": True,
            "download_url": f"/download/{filename}",
            "filename": filename,
            "converter": status["docx_method"]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"转换失败: {str(e)}")

@app.post("/convert-to-pdf")
async def convert_pdf(request: ConvertRequest):
    """将Markdown文本转换为PDF文档"""
    try:
        if not is_pdf_available():
            raise HTTPException(status_code=400, detail="PDF转换需要安装Pandoc和LaTeX引擎")
        
        file_path = convert_to_pdf(request.markdown_text, request.title, request.user_question)
        
        # 返回文件下载链接
        filename = os.path.basename(file_path)
        return {
            "success": True,
            "download_url": f"/download/{filename}",
            "filename": filename,
            "converter": "Pandoc"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"转换失败: {str(e)}")

@app.get("/download/{filename}")
async def download_file(filename: str):
    """下载生成的文档文件"""
    file_path = os.path.join("temp", filename)
    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="文件不存在")
    
    # 根据文件扩展名设置MIME类型
    if filename.endswith('.pdf'):
        media_type = 'application/pdf'
    elif filename.endswith('.docx'):
        media_type = 'application/vnd.openxmlformats-officedocument.wordprocessingml.document'
    else:
        media_type = 'application/octet-stream'
    
    return FileResponse(
        path=file_path,
        filename=filename,
        media_type=media_type
    )

@app.get("/converter-status")
async def converter_status():
    """获取转换器状态信息"""
    return get_converter_status()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            # 接收用户消息
            data = await websocket.receive_text()
            message_data = json.loads(data)
            user_message = message_data.get("message", "")
            
            if not user_message.strip():
                continue
                
            # 获取历史对话
            history = manager.chat_histories.get(websocket, [])
            
            # 发送开始响应标识
            await websocket.send_text(json.dumps({
                "type": "start",
                "message": "AI正在思考..."
            }))
            
            # 收集完整的AI响应
            full_response = ""
            
            # 流式生成响应
            async for chunk in chatbot.stream_response(user_message, history):
                full_response += chunk
                # 发送流式数据
                await websocket.send_text(json.dumps({
                    "type": "stream",
                    "chunk": chunk
                }))
                # 添加小延迟以模拟真实的流式效果
                await asyncio.sleep(0.01)
            
            # 发送结束标识
            await websocket.send_text(json.dumps({
                "type": "end",
                "message": "响应完成"
            }))
            
            # 更新对话历史
            history.append({"role": "user", "content": user_message})
            history.append({"role": "assistant", "content": full_response})
            manager.chat_histories[websocket] = history
            
    except WebSocketDisconnect:
        manager.disconnect(websocket)
    except Exception as e:
        await websocket.send_text(json.dumps({
            "type": "error",
            "message": f"发生错误: {str(e)}"
        }))
        manager.disconnect(websocket)

@app.get("/")
async def get():
    # 获取项目根目录
    root_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
    static_path = os.path.join(root_dir, "static", "index.html")
    return HTMLResponse(content=open(static_path, "r", encoding="utf-8").read())

# 挂载静态文件
try:
    root_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
    static_dir = os.path.join(root_dir, "static")
    app.mount("/static", StaticFiles(directory=static_dir), name="static")
except:
    pass  # 如果static目录不存在，忽略错误

if __name__ == "__main__":
    port = int(os.getenv("PORT", "8000"))
    uvicorn.run(app, host="0.0.0.0", port=port)